#ifndef PCPS_MARKOV_CHAIN_HEADER
#define PCPS_MARKOV_CHAIN_HEADER

#include <src/pcps.h>
#include <src/input.h>
#include <src/hamiltonian.h>
#include <src/wavefunction.h>
#include <src/arena.h>
#include <src/random.h>
#include <src/move.h>
#include <src/coeff_value.h>
#include <src/accumulator.h>

namespace pcps {

//--------------------------------------------------------------------------------
// pcps::MarkovSpecialAction -- Base class for function objects that have the
//                              threads perform special actions during the markov
//                              chain evolution.
//--------------------------------------------------------------------------------

template <class S, class PART, class REF> class MarkovSpecialAction {

  public:

    virtual void action(const int n_affected, int * const affected, S & coeff_ratio,
                        int * const new_config, int * const old_config, pcps::CoeffValueFunc<S, PART, REF> & cvf,
                        pcps::ThreadAccum<S, PART, REF> ** accumulators, const int n_accumulators) = 0;

    virtual ~MarkovSpecialAction() {};

};

//--------------------------------------------------------------------------------
// pcps::MarkovArgPack -- Class to hold the arguments for the markov_chain_detail
//                        function.
//--------------------------------------------------------------------------------

template <class S, class PART, class REF> class MarkovArgPack {

  public:
    const pcps::Input * userinp;
    const pcps::Hamiltonian<S> * ham;
    pcps::Wavefunction<S> * wfn;
    std::vector<boost::shared_ptr<pcps::ProcessAccum<S, PART, REF> > > * accumulators;
    int tid;
    int nthread;
    bool * proceed;
    bool * burn_in;
    bool * master_waiting;
    int * n_slaves_waiting;
    long int * max_iter;
    pthread_mutex_t * cond_mutex;
    pthread_cond_t * cond_slave;
    pthread_cond_t * cond_master;
    pcps::MarkovSpecialAction<S, PART, REF> ** special_action;

};

//--------------------------------------------------------------------------------
// pcps::markov_wait_for_master -- Function in which a slave thread waits for the
//                                 signal to continue from the master thread.
//--------------------------------------------------------------------------------

static void markov_wait_for_master(const int nthread,
                                   const bool & master_waiting,
                                   int & n_slaves_waiting,
                                   pthread_mutex_t * const cond_mutex,
                                   pthread_cond_t * const cond_master,
                                   pthread_cond_t * const cond_slave) {
                            
        pthread_mutex_lock(cond_mutex);
        n_slaves_waiting++;
        if (n_slaves_waiting == nthread && master_waiting)
          pthread_cond_signal(cond_master);
        pthread_cond_wait(cond_slave, cond_mutex);
        pthread_mutex_unlock(cond_mutex);
}

//--------------------------------------------------------------------------------
// pcps::markov_chain_detail -- Iterates one thread's markov chain.
//                              The exact behavior of the markov chain is
//                              determined by the choice of template parameters.
//--------------------------------------------------------------------------------

template <class S, class PART, class REF>
static void * markov_chain_detail(void * arg) {

  // get a reference to the argument pack
  const pcps::MarkovArgPack<S, PART, REF> & arg_pack = *(pcps::MarkovArgPack<S, PART, REF> *)arg;

  // unpack the input arguments
  const pcps::Input & userinp = *arg_pack.userinp;
  const pcps::Hamiltonian<S> & ham = *arg_pack.ham;
  pcps::Wavefunction<S> & wfn = *arg_pack.wfn;
  std::vector<boost::shared_ptr<pcps::ProcessAccum<S, PART, REF> > > & process_accums = *arg_pack.accumulators;
  const int tid = arg_pack.tid;
  const int nthread = arg_pack.nthread;
  const bool & proceed = *arg_pack.proceed;
  const bool & burn_in = *arg_pack.burn_in;
  const bool & master_waiting = *arg_pack.master_waiting;
  int & n_slaves_waiting = *arg_pack.n_slaves_waiting;
  const long int & max_iter = *arg_pack.max_iter;
  pthread_mutex_t * cond_mutex = arg_pack.cond_mutex;
  pthread_cond_t * cond_slave = arg_pack.cond_slave;
  pthread_cond_t * cond_master = arg_pack.cond_master;
  pcps::MarkovSpecialAction<S, PART, REF> ** special_action = arg_pack.special_action;

  // get the number of accumulators
  const int n_accum = process_accums.size();

  // create a memory allocator for the thread (see the Arena class for why this is necessary)
  pcps::Arena arena(userinp.arena_block_size(), userinp.arena_buffer_size());

  // create a random number generator
  pcps::LaggedFibonacci<1279,418> gen(arena);

  // create arrays for the system's configuration
  int * const new_config = arena.allocate_array<int>(userinp.nsites(), 0);
  int * const old_config = arena.allocate_array<int>(userinp.nsites(), 0);

  // declare variables for tracking which sites are affected by a move
  int n_affected = 0;
  int * const affected = arena.allocate_array<int>(4, 0);

  // create object to calculate configuration coefficients
  pcps::CoeffValueFunc<S, PART, REF> cvf(userinp, ham, wfn, arena);

  // create object to propose moves
  pcps::MoveFunc<S, PART, REF> move_func(userinp, ham, wfn, arena);

  // create thread accumulators
  pcps::ThreadAccum<S, PART, REF> ** const thread_accums = arena.allocate_array<pcps::ThreadAccum<S, PART, REF> *>(n_accum);
  for (int i = 0; i < n_accum; i++)
    thread_accums[i] = process_accums.at(i)->create_thread_accum(tid, arena);

  // initialize the system's configuration
  move_func.initialize_configuration(userinp, cvf, new_config, gen);
  for (int i = 0; i < userinp.nsites(); i++)
    old_config[i] = new_config[i];

  // initialize the object that computes configuration coefficients for the chosen initial configuration
  cvf.init(new_config);

  // initialize a variable for the ratio between two configurations' coefficients
  S coeff_ratio = pcps::unity<S>();

  // initialize the iteration number
  long int iter = 0;

  // initialize counter for how many times the current configuration has been sampled
  long int count = 1;

  // wait for instructions from master thread, and complete any special actions that are requested
  pcps::markov_wait_for_master(nthread, master_waiting, n_slaves_waiting, cond_mutex, cond_master, cond_slave);
  while (*special_action != NULL) {
    (**special_action).action(n_affected, affected, coeff_ratio, new_config, old_config, cvf, thread_accums, n_accum);
    pcps::markov_wait_for_master(nthread, master_waiting, n_slaves_waiting, cond_mutex, cond_master, cond_slave);
  }

  // iterate the markov chain until the master thread signals us to stop via the proceed flag
  while (proceed) {

    // make the proposed move and compute the ratio of proposal densities q
    const double q = move_func.propose(n_affected, affected, new_config, gen);

    // calculate the ratio of the configurations' coefficients (new_config's coefficient in numerator)
    coeff_ratio = cvf.get_coeff_ratio(n_affected, affected, new_config, old_config);

    // determine whether or not to accept the move
    const bool accept = ( pcps::uni_01(gen) < q * pcps::square_norm(coeff_ratio) );

    //// print the configuration, the proposal density ratio, and the square of the coefficient value ratio
    //if (tid == 0) {
    //  if (MPI::COMM_WORLD.Get_rank() == 0) {
    //    for (int i = 0; i < userinp.nsites(); i++)
    //      std::cout << old_config[i];
    //    std::cout << boost::format("  %14.4e * %14.4e    %s") % q % pcps::square_norm(coeff_ratio) % (accept ? "accept" : "") << std::endl;
    //  }
    //}

    // increment the iteration number
    iter++;

    // check if we are finished
    assert(max_iter > 0);
    const bool finished = (iter == max_iter);

    // accumulate data for the old configuration
    if (!burn_in)
      for (int i = 0; i < n_accum; i++)
        thread_accums[i]->accumulate(iter, count, accept, finished, coeff_ratio, new_config, old_config, cvf, thread_accums, n_accum);

    // if we are finished iterating, wait for further instructions from the master thread
    if (finished) {

      // wait for instructions from master thread, and complete any special actions that are requested
      pcps::markov_wait_for_master(nthread, master_waiting, n_slaves_waiting, cond_mutex, cond_master, cond_slave);
      while (*special_action != NULL) {
        (**special_action).action(n_affected, affected, coeff_ratio, new_config, old_config, cvf, thread_accums, n_accum);
        pcps::markov_wait_for_master(nthread, master_waiting, n_slaves_waiting, cond_mutex, cond_master, cond_slave);
      }

      // reset the iteration number and sample count
      iter = 0;
      count = 0;

    }

    // update the iteration variables
    if (accept) {
      count = 1;
      cvf.update(n_affected, affected, new_config, old_config);
      for (int i = 0; i < n_affected; i++)
        old_config[affected[i]] = new_config[affected[i]];
      move_func.accept(n_affected, affected, new_config);
    } else {
      count++;
      for (int i = 0; i < n_affected; i++)
        new_config[affected[i]] = old_config[affected[i]];
    }

  }

  // force the thread accumulators to call their destructors (their memory will be deallocated later when arena goes out of scope)
  pthread_mutex_lock(pcps::alloc_mutex());
  for (int i = 0; i < n_accum; i++)
    delete thread_accums[i];
  pthread_mutex_unlock(pcps::alloc_mutex());

  // arena goes out of scope here and the memory it allocated is deallocated
  // (including that for the thread accumulators, although their destructors are not called automatically.  See above.)

}

//--------------------------------------------------------------------------------
// pcps::markov_chain -- A class for iterating markov chains on multiple threads.
//                       The exact behavior of the markov chain is determined
//                       by the choice of template parameters.
//--------------------------------------------------------------------------------

template <class S, class PART, class REF> class MarkovChain {

  private:

    ////////////////////////////////////////////////////////////////////////////
    // member data
    ////////////////////////////////////////////////////////////////////////////
    const int _nthread;
    std::vector<pthread_t> _threads;
    std::vector< pcps::MarkovArgPack<S, PART, REF> > _thread_args;
    bool _proceed;
    bool _burn_in;
    bool _master_waiting;
    int _n_slaves_waiting;
    long int _max_iter;
    pthread_mutex_t _cond_mutex;
    pthread_cond_t _cond_slave;
    pthread_cond_t _cond_master;
    pcps::MarkovSpecialAction<S, PART, REF> * _special_action;

  public:

    ////////////////////////////////////////////////////////////////////////////
    // constructor
    ////////////////////////////////////////////////////////////////////////////
    MarkovChain(const pcps::Input & userinp,
                const pcps::Hamiltonian<S> & ham,
                pcps::Wavefunction<S> & wfn,
                std::vector<boost::shared_ptr<pcps::ProcessAccum<S, PART, REF> > > & accumulators)
      : _nthread(userinp.nthreads()),
        _threads(userinp.nthreads()),
        _thread_args(userinp.nthreads()),
        _proceed(true),
        _burn_in(false),
        _master_waiting(false),
        _n_slaves_waiting(0),
        _max_iter(1),
        _special_action(NULL)
    {

      // initialize mutex and condition variables
      pthread_mutex_init(&_cond_mutex, NULL);
      pthread_cond_init(&_cond_slave, NULL);
      pthread_cond_init(&_cond_master, NULL);

      // create thread attribute object to ensure threads are joinable
      pthread_attr_t attr;
      pthread_attr_init(&attr);
      pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_JOINABLE);

      // initialize the argument packs and start the threads
      for (int i = 0; i < _nthread; i++) {
        _thread_args[i].userinp = &userinp;
        _thread_args[i].ham = &ham;
        _thread_args[i].wfn = &wfn;
        _thread_args[i].accumulators = &accumulators;
        _thread_args[i].tid = i;
        _thread_args[i].nthread = _nthread;
        _thread_args[i].proceed = &_proceed;
        _thread_args[i].burn_in = &_burn_in;
        _thread_args[i].master_waiting = &_master_waiting;
        _thread_args[i].n_slaves_waiting = &_n_slaves_waiting;
        _thread_args[i].max_iter = &_max_iter;
        _thread_args[i].cond_mutex = &_cond_mutex;
        _thread_args[i].cond_slave = &_cond_slave;
        _thread_args[i].cond_master = &_cond_master;
        _thread_args[i].special_action = &_special_action;
        const int rc = pthread_create(&_threads.at(i), &attr, pcps::markov_chain_detail<S, PART, REF>, (void *)&_thread_args[i]);
        if (rc != 0)
          throw pcps::Exception( (boost::format("error in pcps::MarkovChain constructor.  pthread_create failed for thread %i with return code %i")
                                  % i % rc).str() );
      }

      // wait for threads to finish their initialization
      pthread_mutex_lock(&_cond_mutex);
      if (_n_slaves_waiting < _nthread) {
        _master_waiting = true;
        pthread_cond_wait(&_cond_master, &_cond_mutex);
        _master_waiting = false;
      }
      pthread_mutex_unlock(&_cond_mutex);

      // destroy thread attribute object
      pthread_attr_destroy(&attr);

    }

    ////////////////////////////////////////////////////////////////////////////
    // destructor
    ////////////////////////////////////////////////////////////////////////////
    ~MarkovChain() {

      // wake threads and tell them to exit their routines
      pthread_mutex_lock(&_cond_mutex);
      if (_n_slaves_waiting != _nthread)
        throw pcps::Exception( (boost::format("error in pcps::MarkovChain destructor.  Only %i of %i slaves were waiting.")
                                % _n_slaves_waiting % _nthread).str() );
      pthread_cond_broadcast(&_cond_slave);
      _n_slaves_waiting = 0;
      _proceed = false;
      pthread_mutex_unlock(&_cond_mutex);

      // wait for threads to finish
      void * status;
      for (int i = 0; i < _nthread; i++) {
        const int rc = pthread_join(_threads.at(i), &status);
        if (rc != 0)
          throw pcps::Exception( (boost::format("error in pcps::MarkovChain destructor.  pthread_join failed for thread %i with return code %i.")
                                   % i % rc).str() );
      }

      // destroy mutex and condition variables
      pthread_cond_destroy(&_cond_master);
      pthread_cond_destroy(&_cond_slave);
      pthread_mutex_destroy(&_cond_mutex);

    }

    ////////////////////////////////////////////////////////////////////////////
    // function to make the threads do some work
    // (what the threads do depends on the settings of the control variables
    //  and which accumulators the markov chain is handling)
    ////////////////////////////////////////////////////////////////////////////
    void work_threads(const long int sample_length, const bool burn, pcps::MarkovSpecialAction<S, PART, REF> * special_action) {

      pthread_mutex_lock(&_cond_mutex);

      // check that all threads are waiting
      if (_n_slaves_waiting != _nthread)
        throw pcps::Exception( (boost::format("error on entry to pcps::MarkovChain::work_threads.  Only %i of %i slaves were waiting.")
                                % _n_slaves_waiting % _nthread).str() );

      // set control variables, which tell the slave threads what work they should do
      _n_slaves_waiting = 0;
      _max_iter = sample_length;
      _burn_in = burn;
      _special_action = special_action;
      _master_waiting = true;

      // wake threads (they should be waiting inside the markov_wait_for_master function)
      pthread_cond_broadcast(&_cond_slave);

      // wait for threads to finish their work
      pthread_cond_wait(&_cond_master, &_cond_mutex);

      // reset control variables
      _max_iter = 0;
      _burn_in = false;
      _special_action = NULL;
      _master_waiting = false;

      // check that all threads are waiting
      if (_n_slaves_waiting != _nthread)
        throw pcps::Exception( (boost::format("error on exit of pcps::MarkovChain::work_threads.  Only %i of %i slaves were waiting.")
                                % _n_slaves_waiting % _nthread).str() );

      pthread_mutex_unlock(&_cond_mutex);

    }

    ////////////////////////////////////////////////////////////////////////////
    // function to sample the markov chain
    ////////////////////////////////////////////////////////////////////////////
    void sample(const long int sample_length) { this->work_threads(sample_length, false, NULL); }

    ////////////////////////////////////////////////////////////////////////////
    // function to burn in the markov chain
    ////////////////////////////////////////////////////////////////////////////
    void burn_in(const long int sample_length) { this->work_threads(sample_length, true, NULL); }

    ////////////////////////////////////////////////////////////////////////////
    // function object and function for resetting accumulators
    ////////////////////////////////////////////////////////////////////////////
    class ResetAccumulators : public pcps::MarkovSpecialAction<S, PART, REF> {

      public:

        void action(const int n_affected, int * const affected, S & coeff_ratio,
                    int * const new_config, int * const old_config, pcps::CoeffValueFunc<S, PART, REF> & cvf,
                    pcps::ThreadAccum<S, PART, REF> ** accumulators, const int n_accumulators) {

          for (int i = 0; i < n_accumulators; i++)
            accumulators[i]->reset();

        }

    };

    void reset_accumulators() { ResetAccumulators ra; this->work_threads(0, false, &ra); }

    ////////////////////////////////////////////////////////////////////////////
    // function object and function for resetting the coefficient value function
    ////////////////////////////////////////////////////////////////////////////
    class ResetCVF : public pcps::MarkovSpecialAction<S, PART, REF> {

      public:

        void action(const int n_affected, int * const affected, S & coeff_ratio,
                    int * const new_config, int * const old_config, pcps::CoeffValueFunc<S, PART, REF> & cvf,
                    pcps::ThreadAccum<S, PART, REF> ** accumulators, const int n_accumulators) {

          cvf.init(old_config);

        }

    };

    void reset_cvf() { ResetCVF rcvf; this->work_threads(0, false, &rcvf); }

    ////////////////////////////////////////////////////////////////////////////
    // function object and function for computing the wavefunction update
    ////////////////////////////////////////////////////////////////////////////
    class WavefunctionUpdate : public pcps::MarkovSpecialAction<S, PART, REF> {

      public:

        void action(const int n_affected, int * const affected, S & coeff_ratio,
                    int * const new_config, int * const old_config, pcps::CoeffValueFunc<S, PART, REF> & cvf,
                    pcps::ThreadAccum<S, PART, REF> ** accumulators, const int n_accumulators) {

          for (int i = 0; i < n_accumulators; i++)
            accumulators[i]->compute_update();

        }

    };

    void compute_update() { WavefunctionUpdate wu; this->work_threads(0, false, &wu); }

};

} // end namespace pcps

#endif
